package com.wispx.gateway.core.filter.flow.algorithm;

import com.wispx.gateway.core.context.GatewayContext;
import com.wispx.gateway.core.filter.flow.RateLimiter;
import io.netty.channel.EventLoopGroup;

import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

public class LeakyBucketRateLimiter implements RateLimiter {

    /**
     * 漏桶的容量
     */
    private int capacity;

    /**
     * 漏水的速率，单位 ms/个。 eg：200ms/个
     */
    private int leakyWaterRate;

    /**
     * 漏桶中的水量
     */
    private AtomicInteger currentWaterInBucket;

    /**
     * 请求队列
     */
    private Queue<GatewayContext> requestQueue;

    public LeakyBucketRateLimiter(int capacity, int leakyWaterRate, EventLoopGroup executors) {
        this.capacity = capacity;
        this.leakyWaterRate = leakyWaterRate;
        this.requestQueue = new ConcurrentLinkedQueue<>();
        this.currentWaterInBucket = new AtomicInteger(0);
        startHandleRequestTask(executors);// 开启定时任务
    }

    @Override
    public void tryConsume(GatewayContext context) {
        if (currentWaterInBucket.get() < capacity) {
            currentWaterInBucket.incrementAndGet();
            requestQueue.offer(context);
        } else {
            throw new RuntimeException("流量到达限流上限，该请求将被丢弃");
        }
    }

    private void startHandleRequestTask(EventLoopGroup executors) {
        // 使用 Netty 的定时任务来按频率漏出请求
        executors.scheduleAtFixedRate(() -> {
            GatewayContext gatewayContext = requestQueue.poll();
            if (gatewayContext != null) {
                // 重新提交请求到 Netty 事件循环
                gatewayContext.getNettyContext().channel().eventLoop().execute(() -> {
                    currentWaterInBucket.incrementAndGet();
                    gatewayContext.doFilter();
                });
            }
        }, leakyWaterRate, leakyWaterRate, TimeUnit.MILLISECONDS);
    }
}
